/*
 * AESUtil.java
 * Copyright 2021 Qunhe Tech, all rights reserved.
 * Qunhe PROPRIETARY/CONFIDENTIAL, any form of usage is subject to approval.
 */

package top.bbman.boot.auth.component;

import cn.hutool.core.codec.Base64;
import cn.hutool.core.util.StrUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import top.bbman.boot.auth.config.LogInConfig;
import top.bbman.boot.repository.constant.PeaBeanName;

import javax.crypto.Cipher;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

/**
 * @author 祥霸
 * @version 1.0.0
 */
@Component(PeaBeanName.Auth.AES)
public class AESComponent {

    @Autowired
    private LogInConfig logInConfig;

    private static final ThreadLocal<Map<String, Cipher>> CIPHER_THREAD_LOCAL_EN = new ThreadLocal<>();

    private static final ThreadLocal<Map<String, Cipher>> CIPHER_THREAD_LOCAL_DE = new ThreadLocal<>();

    /**
     * 加密
     *
     * @param sSrc 加密内容
     * @param aesKey 密钥
     * @return str
     */
    private static String encrypt(String sSrc, String aesKey) throws Exception {
        if (StrUtil.isBlank(sSrc) || StrUtil.isBlank(aesKey)) {
            return null;
        }
        Map<String, Cipher> cipherMap = CIPHER_THREAD_LOCAL_EN.get();

        if (cipherMap == null) {
            cipherMap = new HashMap<>(10);
            CIPHER_THREAD_LOCAL_EN.set(cipherMap);
        }
        Cipher aesCipher = cipherMap.get(aesKey);
        if (aesCipher == null) {
            byte[] raw = aesKey.getBytes(StandardCharsets.UTF_8);
            SecretKeySpec skeySpec = new SecretKeySpec(raw, "AES");
            aesCipher = Cipher.getInstance("AES");
            aesCipher.init(Cipher.ENCRYPT_MODE, skeySpec);
            cipherMap.put(aesKey, aesCipher);
        }

        return Base64.encode(aesCipher.doFinal(sSrc.getBytes(StandardCharsets.UTF_8)));
    }

    /**
     * 解密
     *
     * @param sSrc 内容
     * @param aesKey 密钥
     * @return str
     */
    private static String decrypt(String sSrc, String aesKey) throws Exception  {
        if (StrUtil.isBlank(sSrc) || StrUtil.isBlank(aesKey)) {
            return "";
        }

        Map<String, Cipher> cipherMap = CIPHER_THREAD_LOCAL_DE.get();
        if (cipherMap == null) {
            cipherMap = new HashMap<>(10);
            CIPHER_THREAD_LOCAL_DE.set(cipherMap);
        }
        Cipher aesCipher = cipherMap.get(aesKey);
        if (aesCipher == null) {
            byte[] raw = aesKey.getBytes(StandardCharsets.UTF_8);
            SecretKeySpec skeySpec = new SecretKeySpec(raw, "AES");
            aesCipher = Cipher.getInstance("AES");
            aesCipher.init(Cipher.DECRYPT_MODE, skeySpec);
            cipherMap.put(aesKey, aesCipher);
        }
        return new String(aesCipher.doFinal(Base64.decode(sSrc)), StandardCharsets.UTF_8);
    }

    /**
     * 使用默认密钥加密
     * @param sSrc 加密内容
     * @return str
     */
    public String encrypt(String sSrc) throws Exception {
        return encrypt(sSrc, this.logInConfig.getAesKey());
    }

    /**
     * 使用默认密钥解密
     * @param sSrc 加密内容
     * @return str
     */
    public String decrypt(String sSrc) throws Exception {
        return decrypt(sSrc, this.logInConfig.getAesKey());
    }
}
